#!/bin/bash

# 脚本名称：disk_classifier.sh

# 功能说明：
# 本脚本用于分类和列出不同类型的存储设备信息。

# 主要流程包括：
# 1. 安装必要的依赖工具（如 lsblk 等）。
# 2. 获取所有存储设备信息。
# 3. 获取可移动磁盘（U盘、移动硬盘）。
# 4. 获取固态硬盘 (SSD)。
# 5. 获取机械硬盘 (HDD)。
# 6. 获取 NVMe 设备。
# 7. 获取光驱设备 (CD-ROM)。
# 8. 获取 Loop 设备（挂载镜像）。
# 9. 获取 NBD 设备（QEMU / 远程块存储）。

# 注意事项：
# - 需要 root 权限运行部分命令。
# - 在执行磁盘操作前，建议备份重要数据，以防万一。

# 获取所有存储设备信息
list_all_disks() {
    lsblk -o NAME,MAJ:MIN,RM,SIZE,FSTYPE,UUID,RO,TYPE,MOUNTPOINT,VENDOR,HOTPLUG -e7
}

# 获取可移动磁盘（U盘、移动硬盘）
list_removable_disks() {
    local removable_disks=($(grep -Hv ^0$ /sys/block/*/removable | sed s/removable:.*$/device\\/uevent/ |
      xargs grep -H ^DRIVER=sd | sed s/device.uevent.*$/size/ | xargs grep -Hv ^0$ |
      cut -d / -f 4))

    for i in "${!removable_disks[@]}"; do
        removable_disks[$i]="/dev/${removable_disks[$i]}"
    done

    echo "🔹 可移动磁盘 (U盘 / 移动硬盘)："
    printf "%s\n" "${removable_disks[@]}"
}

# 获取固态硬盘 (SSD)
list_ssd_disks() {
    local ssd_disks=($(lsblk -d -o NAME,ROTA | awk '$2 == 0 {print "/dev/" $1}'))
    
    echo "⚡ 固态硬盘 (SSD)："
    printf "%s\n" "${ssd_disks[@]}"
}

# 获取机械硬盘 (HDD)
list_hdd_disks() {
    local hdd_disks=($(lsblk -d -o NAME,ROTA | awk '$2 == 1 {print "/dev/" $1}'))
    
    echo "💾 机械硬盘 (HDD)："
    printf "%s\n" "${hdd_disks[@]}"
}

# 获取 NVMe 设备
list_nvme_disks() {
    local nvme_disks=($(lsblk -d -o NAME | grep -E '^nvme'))
    
    echo "🚀 NVMe 硬盘："
    for disk in "${nvme_disks[@]}"; do
        echo "/dev/$disk"
    done
}

# 获取光驱设备 (CD-ROM)
list_cdrom_disks() {
    local cdrom_disks=($(lsblk -d -o NAME,TYPE | awk '$2=="rom" {print "/dev/"$1}'))
    
    echo "📀 光盘驱动器 (CD-ROM)："
    printf "%s\n" "${cdrom_disks[@]}"
}

# 获取 Loop 设备（挂载镜像）
list_loop_devices() {
    local loop_disks=($(lsblk -d -o NAME,TYPE | awk '$2=="loop" {print "/dev/"$1}'))

    echo "🔁 Loop 设备 (ISO/IMG/QCOW2 挂载)："
    printf "%s\n" "${loop_disks[@]}"
}

# 获取 NBD 设备（QEMU / 远程块存储）
list_nbd_devices() {
    local nbd_disks=($(lsblk -d -o NAME | grep -E '^nbd'))
    
    echo "🌐 NBD 设备 (QEMU / 远程块存储)："
    for disk in "${nbd_disks[@]}"; do
        echo "/dev/$disk"
    done
}

# 主函数
main() {
    echo "📊 **磁盘分类信息**"
    echo "======================"
    list_all_disks
    echo "======================"
    list_removable_disks
    echo "======================"
    list_ssd_disks
    echo "======================"
    list_hdd_disks
    echo "======================"
    list_nvme_disks
    echo "======================"
    list_cdrom_disks
    echo "======================"
    list_loop_devices
    echo "======================"
    list_nbd_devices
}

# 运行主函数
main
